# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# --------------------------------------------------------
# References:
# ELECTRA https://github.com/google-research/electra
# BEiT: https://github.com/microsoft/unilm/tree/master/beit
# --------------------------------------------------------

import json


def param_groups_lrd(model, weight_decay=0.05, no_weight_decay_list=[], layer_decay=.75):
    """
    计算用于层级学习率衰减的参数组。
    该方法依据 BEiT 的优化器设计（来自微软的 BEiT 模型），实现了每层不同的学习率衰减策略。

    参数：
        model (nn.Module): 需要优化的模型。
        weight_decay (float, optional): 权重衰减的值（L2正则化）。默认为0.05。
        no_weight_decay_list (list, optional): 不使用权重衰减的参数列表。默认为空列表。
        layer_decay (float, optional): 每层的学习率衰减系数。默认为0.75。

    返回：
        list: 包含各个参数组的列表，每个参数组都有相应的学习率衰减和权重衰减。
    """
    param_group_names = {}  # 存储每个参数组的名称及其设置
    param_groups = {}  # 存储实际的参数组信息

    num_layers = len(model.blocks) + 1  # 获取模型中的层数（假设 blocks 是一个层列表）

    # 为每一层计算衰减系数
    layer_scales = list(layer_decay ** (num_layers - i) for i in range(num_layers + 1))

    # 遍历模型中的所有参数
    for n, p in model.named_parameters():
        if not p.requires_grad:
            continue  # 跳过不需要梯度的参数

        # 对于一维参数和在 no_weight_decay_list 中列出的参数，不使用权重衰减
        if p.ndim == 1 or n in no_weight_decay_list:
            g_decay = "no_decay"  # 无衰减
            this_decay = 0.  # 无权重衰减
        else:
            g_decay = "decay"  # 使用衰减
            this_decay = weight_decay  # 权重衰减系数

        # 获取该参数所在层的层级
        layer_id = get_layer_id_for_vit(n, num_layers)

        # 构造参数组的名称
        group_name = "layer_%d_%s" % (layer_id, g_decay)

        # 如果该层的参数组名称尚未在param_group_names中出现，创建新的参数组
        if group_name not in param_group_names:
            this_scale = layer_scales[layer_id]  # 获取该层的学习率衰减系数

            # 初始化该层的参数组
            param_group_names[group_name] = {
                "lr_scale": this_scale,
                "weight_decay": this_decay,
                "params": [],  # 存储该组的所有参数
            }
            param_groups[group_name] = {
                "lr_scale": this_scale,
                "weight_decay": this_decay,
                "params": [],  # 存储该组的所有参数
            }

        # 将当前参数添加到对应的参数组
        param_group_names[group_name]["params"].append(n)
        param_groups[group_name]["params"].append(p)

    # 返回所有的参数组
    return list(param_groups.values())


def get_layer_id_for_vit(name, num_layers):
    """
    根据参数名称为其分配层ID。
    该方法用于识别在视觉变换器（ViT）模型中各个参数所对应的层数。

    参数：
        name (str): 参数的名称。
        num_layers (int): 总层数，用于确定最大层ID。

    返回：
        int: 参数所对应的层级ID。
    """
    # 如果参数是 'cls_token' 或 'pos_embed'（即与位置编码相关），它们属于第0层
    if name in ['cls_token', 'pos_embed']:
        return 0
    # 如果参数是与 Patch Embedding 相关的（即第一层），它也属于第0层
    elif name.startswith('patch_embed'):
        return 0
    # 如果参数是 'blocks' 层中的参数，则返回其所在的块编号（例如块0，块1）
    elif name.startswith('blocks'):
        return int(name.split('.')[1]) + 1  # 获取块编号并加1（因为第一层为patch_embed）
    # 如果参数不属于上述情况，则它属于最后一层
    else:
        return num_layers
